import os
import numpy as np
from tensorflow import keras
from .config import *
from .model import build_model 
from .data_loader import create_dataset, day_block_shuffle

import tensorflow as tf
from tensorflow.keras.optimizers import Adam
import torch.nn as nn

def train_model(data_path, times):
    # 数据准备
    indices = day_block_shuffle(times)
    split = int(len(indices) * (1-config.VAL_RATIO))
    train_indices, val_indices = indices[:split], indices[split:]
    
    # 创建数据集
    train_ds = create_dataset(data_path, train_indices, config.BATCH_SIZE)
    val_ds = create_dataset(data_path, val_indices, config.BATCH_SIZE)
    
    # 模型构建
    model = build_model()
    #model = Image2PV()
    model.compile(optimizer=keras.optimizers.Adam(config.LEARNING_RATE), loss='mse')
    
    # 回调函数
    callbacks = [
        keras.callbacks.EarlyStopping(patience=config.PATIENCE, restore_best_weights=True),
        keras.callbacks.ModelCheckpoint(
            os.path.join(config.output_folder, 'best_model.h5'),
            save_best_only=True
        )
    ]
    
    # 训练模型
    history = model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=config.EPOCHS,
        callbacks=callbacks
    )
    
    return history, model
